import os
import pandas as pd
from pathlib import Path
import dolphindb as ddb
import numpy as np
import win32com.client
import tempfile
import time

class SHFEDataImporter:
    def __init__(self):
        self.source_path = Path("F:/资料/历史数据/上期所")
        self.session = None
        self.upserter = None
        self.excel = None
        
    def connect_db(self):
        """连接DolphinDB"""
        self.session = ddb.session()
        self.session.connect("localhost", 8848, "admin", "123456")
        # 使用 TableUpserter 替代 TableAppender
        self.upserter = ddb.TableUpserter(
            dbPath="dfs://futures", 
            tableName="daily_all", 
            ddbSession=self.session,
            keyColNames=["date", "code", "exchange_code"]  # 指定键值列，用于判断重复
        )
        
    def _read_excel(self, file_path):
        """使用多种方法尝试读取Excel文件"""
        try:
            # 方法1: 使用xlrd读取
            try:
                df = pd.read_excel(file_path, header=2, engine='xlrd')
                print("使用xlrd成功读取文件")
                return df
            except Exception as e1:
                print(f"xlrd读取失败: {str(e1)}")

            # 方法2: 使用openpyxl读取
            try:
                df = pd.read_excel(file_path, header=2, engine='openpyxl')
                print("使用openpyxl成功读取文件")
                return df
            except Exception as e2:
                print(f"openpyxl读取失败: {str(e2)}")

            # 方法3: 使用odf读取
            try:
                import odf
                from odf.opendocument import load
                from odf.table import Table, TableRow, TableCell
                
                doc = load(str(file_path))
                table = doc.spreadsheet.getElementsByType(Table)[0]
                data = []
                for row in table.getElementsByType(TableRow):
                    row_data = []
                    for cell in row.getElementsByType(TableCell):
                        row_data.append(cell.getAttribute("value") or "")
                    data.append(row_data)
                df = pd.DataFrame(data[3:], columns=data[2])
                print("使用odf成功读取文件")
                return df
            except Exception as e3:
                print(f"odf读取失败: {str(e3)}")

            # 方法4: 使用pyxlsb读取
            try:
                df = pd.read_excel(file_path, header=2, engine='pyxlsb')
                print("使用pyxlsb成功读取文件")
                return df
            except Exception as e4:
                print(f"pyxlsb读取失败: {str(e4)}")

            print("所有读取方法都失败")
            return None

        except Exception as e:
            print(f"读取Excel文件失败: {str(e)}")
            return None
            
    def is_option_code(self, code):
        """判断是否为期权合约"""
        try:
            code = str(code).upper()
            # 上期所期权合约特征：
            # 1. 包含 'P' 或 'C' 字符（看跌或看涨期权）
            # 2. 期权代码通常较长
            option_indicators = ['P', 'C']
            return any(indicator in code for indicator in option_indicators) and len(code) > 6
        except:
            return False

    def process_file(self, file_path, year):
        """处理单个文件"""
        try:
            # 读取Excel文件
            df = self._read_excel(file_path)
            if df is None:
                return None

            # 根据年份决定数据起始行
            if int(year) < 2020:
                data_start_row = 0
            else:
                data_start_row = 1
            
            # 只保留数据行
            df = df.iloc[data_start_row:]
            
            # 处理合并单元格：向下填充合约列
            df['合约'] = df['合约'].ffill()
            
            # 过滤无效行：
            # 1. 删除所有列都为空的行
            df = df.dropna(how='all')
            
            # 2. 处理日期列名可能的不同情况
            date_columns = ['日期', '交易日期']
            date_col = None
            for col in date_columns:
                if col in df.columns:
                    date_col = col
                    break
                
            if date_col is None:
                print(f"未找到日期列，可用列名: {df.columns.tolist()}")
                return None
            
            # 删除日期为空的行
            df = df.dropna(subset=[date_col])
            
            # 3. 过滤掉注释行（通常以"注："开头）
            df = df[~df['合约'].str.contains('注：', na=False)]
            
            # 4. 确保必要的数值列都有值
            numeric_columns = ['开盘价', '最高价', '最低价', '收盘价']
            df = df.dropna(subset=numeric_columns, how='all')
            
            # 5. 转换成交量为数值类型
            try:
                df['成交量'] = pd.to_numeric(df['成交量'], errors='coerce')
                df = df.dropna(subset=['成交量'])
                df = df[df['成交量'] > 0]
            except Exception as e:
                print(f"转换成交量失败: {str(e)}")
                return None
            
            # 6. 转换其他数值列
            numeric_cols = ['前收盘', '前结算', '开盘价', '最高价', '最低价', '收盘价', '结算价']
            for col in numeric_cols:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')
            
            # 7. 处理成交金额列名可能的不同情况
            amount_columns = ['成交金额', '成交金额(万元)']
            amount_col = None
            for col in amount_columns:
                if col in df.columns:
                    amount_col = col
                    df[col] = pd.to_numeric(df[col], errors='coerce')
                    # 如果是万元，转换为元
                    if '万元' in col:
                        df[col] = df[col] * 10000
                    break
                
            if amount_col is None:
                print(f"未找到成交金额列，可用列名: {df.columns.tolist()}")
                return None
            
            # 重命名列
            rename_dict = {
                date_col: 'date',
                '合约': 'code',
                '最高价': 'high',
                '最低价': 'low',
                '开盘价': 'open',
                '收盘价': 'close',
                '前收盘': 'pre_close',
                '成交量': 'volume',
                amount_col: 'amount'
            }
            
            # 只保留需要的列
            needed_columns = list(rename_dict.keys())
            existing_columns = [col for col in needed_columns if col in df.columns]
            df = df[existing_columns]
            
            df = df.rename(columns=rename_dict)
            
            # 处理日期格式
            df['date'] = pd.to_datetime(df['date'], format='%Y%m%d')
            
            # 添加交易所信息
            df['exchange_code'] = 'SHFE'
            df['exchange_name'] = '上期所'
            
            # 提取品种名称
            df['name'] = df['code'].apply(lambda x: self._get_product_name(x))
            
            # 过滤期权合约
            df = df[~df['code'].apply(self.is_option_code)]
            
            # 选择需要的列
            columns = ['date', 'code', 'name', 'exchange_code', 'exchange_name', 
                      'high', 'low', 'open', 'close', 'pre_close', 'volume', 'amount']
            
            return df[columns]
            
        except Exception as e:
            print(f"处理文件失败 {file_path}: {str(e)}")
            print(f"错误类型: {type(e).__name__}")
            import traceback
            print(f"详细错误信息: {traceback.format_exc()}")
            return None
            
    def _get_product_name(self, code):
        """从合约代码提取品种名称"""
        try:
            # 确保code是字符串
            code = str(code)
            
            # 上期所品种名称映射
            product_names = {
                'cu': '铜',
                'al': '铝',
                'zn': '锌',
                'pb': '铅',
                'ni': '镍',
                'sn': '锡',
                'au': '黄金',
                'ag': '白银',
                'rb': '螺纹钢',
                'wr': '线材',
                'hc': '热轧卷板',
                'fu': '燃料油',
                'bu': '沥青',
                'ru': '天然橡胶'
            }
            
            # 提取品种代码（前2个字符）
            product_code = ''.join(filter(str.isalpha, code[:2].lower()))
            return product_names.get(product_code, product_code)
        except:
            return code
            
    def import_to_db(self, df):
        """使用TableUpserter导入数据，自动处理重复数据"""
        try:
            # 直接使用upserter的upsert方法
            self.upserter.upsert(df)
            print(f"成功处理 {len(df)} 条数据")
        except Exception as e:
            print(f"导入数据失败: {str(e)}")
            # 如果是连接断开，尝试重连
            try:
                print("尝试重新连接...")
                self.connect_db()
                self.upserter.upsert(df)
                print(f"重连后成功处理 {len(df)} 条数据")
            except Exception as e2:
                print(f"重试失败: {str(e2)}")

    def process_all_files(self):
        """处理所有文件"""
        try:
            self.connect_db()
            
            # 遍历所有年份目录
            for year_dir in self.source_path.iterdir():
                if not year_dir.is_dir():
                    continue
                    
                year = year_dir.name
                print(f"\n处理 {year} 年数据...")
                
                # 处理该年份下的所有文件
                for file_path in year_dir.glob('*.xls*'):  # 同时匹配.xls和.xlsx
                    print(f"处理文件: {file_path.name}")
                    df = self.process_file(file_path, year)
                    
                    if df is not None and not df.empty:
                        print(f"处理得到 {len(df)} 行数据")
                        self.import_to_db(df)
                        
        except Exception as e:
            print(f"处理过程出错: {str(e)}")
            raise
        finally:
            if self.session:
                self.session.close()

    def __del__(self):
        """析构函数，确保Excel应用被关闭"""
        if self.excel:
            self.excel.Quit()

def main():
    """主函数"""
    importer = SHFEDataImporter()
    importer.process_all_files()

if __name__ == "__main__":
    main() 